# 指定文件编码为UTF-8
# coding: utf-8

"""
mnist.py 功能说明：
1. 实现MNIST数据集的下载、加载和预处理功能
2. 提供数据归一化、one-hot编码和维度调整选项
3. 自动缓存处理后的数据到pickle文件
4. 支持训练集(60,000样本)和测试集(10,000样本)的加载

主要功能：
- 自动下载MNIST原始数据集
- 将二进制数据转换为NumPy数组
- 数据归一化(0-1范围)
- 标签one-hot编码
- 图像展平或保持原始维度(1x28x28)
- 数据缓存到pickle文件提高后续加载速度

数据格式：
- 图像数据：默认形状为(样本数, 784)的浮点数组(0-1范围)
- 标签数据：整数数组(0-9)或one-hot编码数组

使用示例：
(train_img, train_label), (test_img, test_label) = load_mnist(
    normalize=True, flatten=True, one_hot_label=False)
"""

try:
    import urllib.request  # Python 3.x的URL请求库
except ImportError:
    raise ImportError('You should use Python 3.x')  # 需要Python 3.x环境
import os.path  # 路径操作
import gzip  # 解压gzip文件
import pickle  # 数据序列化
import os  # 操作系统接口
import numpy as np  # 数值计算库

# MNIST数据集的下载URL(镜像站点)
url_base = 'https://ossci-datasets.s3.amazonaws.com/mnist/'
# MNIST数据集的四个关键文件
key_file = {
    'train_img': 'train-images-idx3-ubyte.gz',  # 训练图像
    'train_label': 'train-labels-idx1-ubyte.gz',  # 训练标签
    'test_img': 't10k-images-idx3-ubyte.gz',  # 测试图像
    'test_label': 't10k-labels-idx1-ubyte.gz'  # 测试标签
}

# 数据集保存目录和文件名
dataset_dir = os.path.dirname(os.path.abspath(__file__))
save_file = dataset_dir + "/mnist.pkl"  # 预处理后的数据保存路径

# 数据集参数
train_num = 60000  # 训练样本数
test_num = 10000  # 测试样本数
img_dim = (1, 28, 28)  # 图像原始维度(通道,高,宽)
img_size = 784  # 展平后的图像大小(28*28)

def _download(file_name):
    """下载MNIST数据集文件

    参数:
        file_name: 要下载的文件名
    """
    file_path = dataset_dir + "/" + file_name

    # 如果文件已存在则跳过下载
    if os.path.exists(file_path):
        return

    print("Downloading " + file_name + " ... ")
    urllib.request.urlretrieve(url_base + file_name, file_path)  # 下载文件
    print("Done")

def download_mnist():
    """下载MNIST数据集的所有文件"""
    for v in key_file.values():
       _download(v)

def _load_label(file_name):
    """加载并转换标签文件

    参数:
        file_name: 标签文件名
    返回:
        NumPy数组形式的标签数据
    """
    file_path = dataset_dir + "/" + file_name

    print("Converting " + file_name + " to NumPy Array ...")
    with gzip.open(file_path, 'rb') as f:  # 解压gzip文件
            # 从二进制数据读取标签(跳过前8字节头信息)
            labels = np.frombuffer(f.read(), np.uint8, offset=8)
    print("Done")

    return labels

def _load_img(file_name):
    """加载并转换图像文件

    参数:
        file_name: 图像文件名
    返回:
        NumPy数组形式的图像数据
    """
    file_path = dataset_dir + "/" + file_name

    print("Converting " + file_name + " to NumPy Array ...")
    with gzip.open(file_path, 'rb') as f:  # 解压gzip文件
            # 从二进制数据读取图像(跳过前16字节头信息)
            data = np.frombuffer(f.read(), np.uint8, offset=16)
    data = data.reshape(-1, img_size)  # 重塑为(样本数, 784)形状
    print("Done")

    return data

def _convert_numpy():
    """将MNIST数据集转换为NumPy数组格式

    返回:
        包含训练和测试数据的字典
    """
    dataset = {}
    dataset['train_img'] = _load_img(key_file['train_img'])  # 加载训练图像
    dataset['train_label'] = _load_label(key_file['train_label'])  # 加载训练标签
    dataset['test_img'] = _load_img(key_file['test_img'])  # 加载测试图像
    dataset['test_label'] = _load_label(key_file['test_label'])  # 加载测试标签

    return dataset

def init_mnist():
    """初始化MNIST数据集(下载并转换为pickle文件)"""
    download_mnist()  # 下载原始数据
    dataset = _convert_numpy()  # 转换为NumPy数组
    print("Creating pickle file ...")
    with open(save_file, 'wb') as f:  # 写入pickle文件
        pickle.dump(dataset, f, -1)  # 使用最高协议序列化
    print("Done!")

def _change_one_hot_label(X):
    """将标签转换为one-hot编码

    参数:
        X: 原始标签数组
    返回:
        one-hot编码后的数组
    """
    T = np.zeros((X.size, 10))  # 创建全零矩阵(样本数×10)
    for idx, row in enumerate(T):
        row[X[idx]] = 1  # 对应类别位置设为1

    return T


def load_mnist(normalize=True, flatten=True, one_hot_label=False):
    """加载MNIST数据集

    参数:
        normalize: 是否将像素值归一化到0.0~1.0
        flatten: 是否将图像展平为一维数组
        one_hot_label: 是否将标签转换为one-hot编码
    返回:
        (训练图像, 训练标签), (测试图像, 测试标签)
    """
    if not os.path.exists(save_file):  # 如果pickle文件不存在
        init_mnist()  # 初始化数据集

    with open(save_file, 'rb') as f:  # 读取pickle文件
        dataset = pickle.load(f)  # 加载数据

    if normalize:  # 归一化处理
        for key in ('train_img', 'test_img'):
            dataset[key] = dataset[key].astype(np.float32)  # 转换为浮点数
            dataset[key] /= 255.0  # 归一化到0-1范围

    if one_hot_label:  # one-hot编码转换
        dataset['train_label'] = _change_one_hot_label(dataset['train_label'])
        dataset['test_label'] = _change_one_hot_label(dataset['test_label'])

    if not flatten:  # 保持原始图像维度
         for key in ('train_img', 'test_img'):
            dataset[key] = dataset[key].reshape(-1, 1, 28, 28)  # 重塑为(样本数,1,28,28)

    return (dataset['train_img'], dataset['train_label']), (dataset['test_img'], dataset['test_label'])


if __name__ == '__main__':
    init_mnist()  # 直接运行时初始化数据集
